import codecs
import json
import math
import os
import shutil
import sys
# import time
import shapefile
import cv2
# import imutils
import numpy as np
# import rasterio
# import pyproj
from osgeo import gdal, ogr, osr, gdalconst
import geopandas as gpd
import warnings
from draw import filter_point

python_executable = sys.executable
os.environ['PROJ_LIB'] = os.path.join(os.path.dirname(os.path.dirname(python_executable)), 'share/proj')


def json2shp(in_geojson_file, out_shp_file):
    # 打开GEOJSON文件
    with open(in_geojson_file) as f:
        geojson_data = f.read()
    # 将GEOJSON数据加载到OGR中
    source = ogr.Open(geojson_data)
    layer = source.GetLayer()
    # 获取投影信息
    source_srs = layer.GetSpatialRef()
    # 创建输出SHP文件
    driver = ogr.GetDriverByName('ESRI Shapefile')
    output = driver.CreateDataSource(out_shp_file)
    # 创建新图层
    new_layer = output.CreateLayer('layer_name', source_srs, layer.GetGeomType())
    # 添加字段
    layer_def = layer.GetLayerDefn()
    for i in range(layer_def.GetFieldCount()):
        field_def = layer_def.GetFieldDefn(i)
        new_layer.CreateField(field_def)
    # 将要素从GEOJSON复制到新图层
    for feature in layer:
        geom = feature.GetGeometryRef()
        new_feature = ogr.Feature(new_layer.GetLayerDefn())
        new_feature.SetGeometry(geom)
        for i in range(layer_def.GetFieldCount()):
            field_def = layer_def.GetFieldDefn(i)
            value = feature.GetField(field_def.GetName())
            new_feature.SetField(field_def.GetName(), value)
        new_layer.CreateFeature(new_feature)
    # 保存输出SHP文件
    output.Destroy()
    return


def Shp2JSON(shpfile, outJson, shpencoding):
    """
    这个函数用于将shp文件转换为GeoJSON文件
    :param outJson:
    :param shpfile:
    :param shpencoding:
    :param filename: shp文件对应的文件名（去除文件拓展名）
    :return:
    """

    '''创建shp IO连接'''
    reader = shapefile.Reader(shpfile, encoding=shpencoding)
    # print(len(reader.shapes()))
    '''提取所有field部分内容'''
    fields = reader.fields[1:]
    '''提取所有field的名称'''
    field_names = [field[0] for field in fields]
    '''初始化要素列表'''
    buffer = []
    for sr in reader.shapeRecords():
        '''提取每一个矢量对象对应的属性值'''
        record = sr.record
        '''属性转换为列表'''
        record = [r.decode('gb2312', 'ignore') if isinstance(r, bytes)
                  else r for r in record]
        '''对齐属性与对应数值的键值对'''
        atr = dict(zip(field_names, record))
        '''获取当前矢量对象的类型及矢量信息'''
        geom = sr.shape.__geo_interface__
        '''向要素列表追加新对象'''
        buffer.append(dict(type="Feature", geometry=geom, properties=atr))

    geojson = codecs.open(outJson, 'w', encoding='utf-8')
    geojson.write(json.dumps({"type": "FeatureCollection",
                              "features": buffer}) + '\n')
    return buffer


def create_tif(infile, mask, nodata_out, outfile, reffile):
    # 打开tif文件
    in_ds = gdal.Open(infile, gdal.GA_ReadOnly)  # 输入文件
    ref_ds = gdal.Open(reffile, gdal.GA_ReadOnly)  # 参考文件
    bandCount = in_ds.RasterCount
    ReadBand = in_ds.GetRasterBand(1)
    bandtype = ReadBand.DataType
    nodata = ReadBand.GetNoDataValue()
    # print(nodata)
    # 参考文件与输入文件的的地理仿射变换参数与投影信息
    in_trans = in_ds.GetGeoTransform()
    in_proj = in_ds.GetProjection()
    ref_trans = ref_ds.GetGeoTransform()
    ref_proj = ref_ds.GetProjection()
    # print(in_trans)
    # 输入文件的行列数
    x = in_ds.RasterXSize
    y = in_ds.RasterYSize

    # 创建输出文件
    driver = gdal.GetDriverByName('GTiff')
    output = driver.Create(outfile, x, y, bandCount, bandtype, options=["COMPRESS=LZW"])
    # 设置输出文件地理仿射变换参数与投影
    output.SetGeoTransform(ref_trans)
    output.SetProjection(ref_proj)
    for bandnum in range(bandCount):
        # print(bandnum)
        band = output.GetRasterBand(bandnum + 1)
        if nodata_out is not None:
            band.SetNoDataValue(nodata_out)
        extent_data = in_ds.GetRasterBand(bandnum + 1).ReadAsArray()
        if mask is not None:
            extent_data[mask] = 0
        band.WriteArray(extent_data)
    output = None
    in_ds = None
    ref_ds = None


def pixel_geo_register(infile, outfile, reffile):
    """
    infile:输入文件
    outfile:输出文件
    reffile:参考文件
    methods:重采样方法
            gdalconst.GRA_NearestNeighbour：near
            gdalconst.GRA_Bilinear:bilinear
            gdalconst.GRA_Cubic:cubic
            gdalconst.GRA_CubicSpline:cubicspline
            gdalconst.GRA_Lanczos:lanczos
            gdalconst.GRA_Average:average
            gdalconst.GRA_Mode:mode
    """
    # 打开tif文件
    in_ds = gdal.Open(infile, gdal.GA_ReadOnly)  # 输入文件
    ref_ds = gdal.Open(reffile, gdal.GA_ReadOnly)  # 参考文件
    bandCount = in_ds.RasterCount
    ReadBand = in_ds.GetRasterBand(1)
    bandtype = ReadBand.DataType
    nodata = ReadBand.GetNoDataValue()
    # print(nodata)
    # 参考文件与输入文件的的地理仿射变换参数与投影信息
    in_trans = in_ds.GetGeoTransform()
    in_proj = in_ds.GetProjection()
    ref_trans = ref_ds.GetGeoTransform()
    ref_proj = ref_ds.GetProjection()
    # print(in_trans)
    # 输入文件的行列数
    x = ref_ds.RasterXSize
    y = ref_ds.RasterYSize

    # 创建输出文件
    driver = gdal.GetDriverByName('GTiff')
    output = driver.Create(outfile, x, y, bandCount, bandtype, options=["COMPRESS=LZW"])
    # 设置输出文件地理仿射变换参数与投影
    output.SetGeoTransform(ref_trans)
    output.SetProjection(ref_proj)

    # 重投影，插值方法为双线性内插法
    gdal.ReprojectImage(in_ds, output, in_proj, ref_proj, gdalconst.GRA_NearestNeighbour)

    # 关闭数据集与driver
    in_ds = None
    ref_ds = None
    driver = None
    output = None
    ReadBand = None


def draw_result(raster_path, vector_path, field_name, out_path):
    raster_ds = gdal.Open(raster_path, gdal.GA_Update)
    rows = raster_ds.RasterYSize
    cols = raster_ds.RasterXSize
    bands = raster_ds.RasterCount
    projection = raster_ds.GetProjection()
    geotransform = raster_ds.GetGeoTransform()
    image = np.zeros((rows, cols, bands), dtype=np.uint8)
    for i in range(bands):
        band = raster_ds.GetRasterBand(i + 1)
        band_array = band.ReadAsArray()
        image[:, :, i] = band_array
    vector_ds = ogr.Open(vector_path)
    layer = vector_ds.GetLayer()
    unique_values = set()
    for feature in layer:
        value = feature.GetField(field_name)
        unique_values.add(value)
    colors = [(0, 0, 255), (255, 0, 0), (0, 255, 0), (0, 255, 255), (255, 255, 0), (255, 0, 255)]
    for index, field_value in enumerate(unique_values):
        if field_value == 0: continue
        layer.SetAttributeFilter(f"{field_name} = '{field_value}'")
        mem_driver = gdal.GetDriverByName("MEM")
        mem_raster_ds = mem_driver.Create("", cols, rows, 1, gdal.GDT_Byte)
        mem_raster_ds.SetProjection(projection)
        mem_raster_ds.SetGeoTransform(geotransform)
        gdal.RasterizeLayer(mem_raster_ds, [1], layer, burn_values=[255])
        mem_band = mem_raster_ds.GetRasterBand(1)
        mask = mem_band.ReadAsArray()
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        cv2.drawContours(image, contours, -1, colors[index], 2)
        mem_raster_ds = None
    cv2.imwrite(out_path, image)
    raster_ds = None
    vector_ds = None
    return


def PolygonToRaster_conversion(inshp, field, outImg):

    datasource = ogr.Open(inshp)
    layer = datasource.GetLayer(0)
    extent = layer.GetExtent()

    cols = math.ceil((extent[1] - extent[0]) / 0.000005)
    rows = math.ceil((extent[3] - extent[2]) / 0.000005)

    GeoTransform = (extent[0], 0.000005, 0, extent[3], 0, -0.000005)

    srs = osr.SpatialReference()
    srs.ImportFromEPSG(4326)
    proj = srs.ExportToWkt()
    # print(proj)
    Driver = gdal.GetDriverByName('GTIff')
    Dataset = Driver.Create(outImg, cols, rows, 1, gdal.GDT_Byte, options=["COMPRESS=LZW"])
    Dataset.SetGeoTransform(GeoTransform)
    Dataset.SetProjection(proj)
    maskband = Dataset.GetRasterBand(1)
    maskband.SetNoDataValue(0)

    gdal.RasterizeLayer(Dataset, [1], layer,
                        options=["ALL_TOUCHED=False", "ATTRIBUTE=" + field, "COMPRESS=LZW"])  # 跟shp字段给栅格像元赋值

    datasource.Destroy()
    del Dataset

    return


def RasterToPolygon(tiff, outshp):
    ds = gdal.Open(tiff)
    srcband = ds.GetRasterBand(1)
    maskband = srcband.GetMaskBand()
    gdal.SetConfigOption("GDAL_FILENAME_IS_UTF8", "YES")
    gdal.SetConfigOption("SHAPE_ENCODING", "UTF-8")
    ogr.RegisterAll()
    drv = ogr.GetDriverByName('ESRI Shapefile')
    dst_ds = drv.CreateDataSource(outshp)
    prosrs = osr.SpatialReference(wkt=ds.GetProjection())
    dst_layer = dst_ds.CreateLayer(
        "POLYGON", geom_type=ogr.wkbPolygon, srs=prosrs)
    dst_fieldname = "CLASS"
    fd = ogr.FieldDefn(dst_fieldname, ogr.OFTInteger)
    dst_layer.CreateField(fd)
    gdal.Polygonize(srcband, maskband, dst_layer, 0, [])
    # TODO: temporary: delete ignored values
    dst_ds.Destroy()
    ds = None


def draw_point(tiffile, shapejson, field, draw_point_size, outfile, draw_line):
    """
        tiffile:输入待渲染图像
        shapejson:输入json矢量
        field: json矢量属性字段
        draw_point_size: 点大小
        outfile:输出文件
        draw_line:是否划线
    """
    ref_ds = gdal.Open(tiffile, gdal.GA_ReadOnly)
    # # 参考文件与输入文件的的地理仿射变换参数与投影信息
    ref_trans = ref_ds.GetGeoTransform()
    ref_proj = ref_ds.GetProjection()
    # 输入文件的行列数
    x = ref_ds.RasterXSize
    y = ref_ds.RasterYSize
    dTemp = ref_trans[1] * ref_trans[5] - ref_trans[2] * ref_trans[4]
    del ref_ds

    img = cv2.imread(tiffile)
    pt_xy = None
    with open(shapejson, 'r', encoding='utf-8') as file:
        # 使用json.load()方法解析JSON数据
        data = json.load(file)
    # print(data.get('features'))
    colors = [(50, 50, 150), (1, 1, 255), (255, 1, 1), (1, 255, 1), (1, 255, 255), (255, 255, 1), (255, 1, 255)]
    for fe, feature in enumerate(data.get('features')):
        if field is None:
            point_xy = feature.get('geometry').get('coordinates')
            Xpixel = (ref_trans[5] * (point_xy[0] - ref_trans[0]) - ref_trans[2] * (
                    point_xy[1] - ref_trans[3])) / dTemp + 0.5
            Yline = (ref_trans[1] * (point_xy[1] - ref_trans[3]) - ref_trans[4] * (
                    point_xy[0] - ref_trans[0])) / dTemp + 0.5
            if fe > 0 and draw_line:
                cv2.line(img, pt_xy, (int(Xpixel), int(Yline)), (0, 255, 0), 1)
            cv2.circle(img, (int(Xpixel), int(Yline)), draw_point_size, (0, 0, 255), -1)
            pt_xy = (int(Xpixel), int(Yline))
        else:
            cls = feature.get('properties').get(field)
            if cls is None:
                cls = 0

            point_xy = feature.get('geometry').get('coordinates')
            Xpixel = (ref_trans[5] * (point_xy[0] - ref_trans[0]) - ref_trans[2] * (
                    point_xy[1] - ref_trans[3])) / dTemp + 0.5
            Yline = (ref_trans[1] * (point_xy[1] - ref_trans[3]) - ref_trans[4] * (
                    point_xy[0] - ref_trans[0])) / dTemp + 0.5

            cv2.circle(img, (int(Xpixel), int(Yline)), draw_point_size, colors[int(cls)], -1)
            pt_xy = (int(Xpixel), int(Yline))

    cv2.imwrite(outfile, img)


def image_prepare(image, shape_json, filename, path):
    temp_dir = os.path.join(path, 'temp')
    try:
        os.makedirs(temp_dir)
    except:
        pass

    # 轨迹点可视化
    image_tif = image
    image_draw = os.path.join(temp_dir, 'image_draw_Gaus.png')

    # 渲染点线
    draw_point(image_tif, shape_json, None, 2, image_draw, True)
    # json转shp
    point_shp = os.path.join(temp_dir, filename + 'shape.shp')
    if not os.path.exists(point_shp):
        json2shp(shape_json, point_shp)

    # 点缓冲区
    point_buffer = os.path.join(temp_dir, 'point_buffer.shp')
    point = gpd.read_file(point_shp)
    buffered = point.buffer(0.00015)
    buffered.to_file(point_buffer)
    point_ = gpd.read_file(point_buffer)
    buffered = point_.assign(BUFF_DIST=15)
    buffered.to_file(point_buffer)
    point_bufferToRaster = os.path.join(temp_dir, 'point_bufferToRaster.tif')
    PolygonToRaster_conversion(point_buffer, 'BUFF_DIST', point_bufferToRaster)

    point_bufferToRaster_new = os.path.join(temp_dir, 'point_bufferToRaster_new.tif')
    if not os.path.exists(point_bufferToRaster_new):
        pixel_geo_register(point_bufferToRaster, point_bufferToRaster_new, image)

    extract_data = cv2.imread(point_bufferToRaster_new, 0)
    mask = np.where(extract_data != 15)
    predict_image = os.path.join(temp_dir, filename + '_image_draw_extract.tif')
    create_tif(image_draw, mask, None, predict_image, image)
    return predict_image


# 预测结果可视化
def pre_result(image, shape_json, result_mask, filename, path):
    result_mask = cv2.imdecode(np.fromfile(result_mask,dtype=np.uint8),-1)
    temp_dir = os.path.join(path, 'temp')
    try:
        os.makedirs(temp_dir)
    except:
        pass
    #
    Predict_png_Polygon = os.path.join(temp_dir, filename + '_result_Polygon.png')
    Predict_Polygon_tif = os.path.join(temp_dir, filename + '_Polygon_filter.tif')
    Predict_Polygon_shp = os.path.join(temp_dir, filename + '_result_Polygon.shp')

    cv2.imwrite(Predict_png_Polygon, result_mask)
    create_tif(Predict_png_Polygon, None, 0, Predict_Polygon_tif, image)
    RasterToPolygon(Predict_Polygon_tif, Predict_Polygon_shp)

    # 叠加提取点
    point_shp = os.path.join(temp_dir, filename + 'shape.shp')
    if not os.path.exists(point_shp):
        json2shp(shape_json, point_shp)

    point_shp_df = gpd.read_file(point_shp)
    Polygon_df = gpd.read_file(Predict_Polygon_shp)
    overlay_Polygon = gpd.overlay(df1=point_shp_df,
                                  df2=Polygon_df,
                                  how='union')  # intersection
    result_Polygon = os.path.join(temp_dir, filename + '_result.shp')
    overlay_Polygon.to_file(result_Polygon)

    result_points = os.path.join(path, filename + '_point.shp')
    filter_point(result_Polygon, result_points)
    # geojson结果
    point_geoJSON = os.path.join(path, filename + '_result_points.json')
    try:
        Shp2JSON(result_points, point_geoJSON, 'utf-8')

    except:
        Shp2JSON(result_points, point_geoJSON, 'gbk')

    point_result = os.path.join(path, filename + '_result_points.tif')
    draw_point(image, point_geoJSON, 'CLASS', 3, point_result, False)
    # 删除
    # shutil.rmtree(temp_dir)
    return point_result


if __name__ == '__main__':
    if sys.argv[1] == "image_prepare":
        image_prepare(sys.argv[2],sys.argv[3],sys.argv[4],sys.argv[5])
    elif sys.argv[1] == "pre_result":
        pre_result(sys.argv[2],sys.argv[3],sys.argv[4],sys.argv[5],sys.argv[6])
    # image=r"/data/fengyy/guiji_Data/hmf/dataset/guiji_train_big/000002_1717114774878/image.png" 
    # result_mask =r"/data/fengyy/guiji_Data/out/000002_1717114774878_image_draw_extract.tif"    
    # filename  =  "000002_1717114774878"
    # path=r"/data/fengyy/guiji_Data/out"    
    # pre_result(image, result_mask, filename, path)

